"""
qwen2指令微调模型-推理(预测)脚本
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def predict(messages, model, tokenizer):
    # 使用GPU进行计算
    device = "cuda"
    # 使用tokenizer.apply_chat_template将messages格式化为适合模型输入的字符串text
    text = tokenizer.apply_chat_template(
        messages,
        # 输入是否要分词，这里不需要分词，直接输出文本
        tokenize=False,
        # 在文本末尾添加生成提示
        add_generation_prompt=True
    )
    """
    使用tokenizer将格式化后的text编码为PyTorch张量，并传入模型进行处理。
    return_tensors="pt"表示返回PyTorch张量。
    将张量移动到指定设备（GPU）上。
    """
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    # 调用model.generate生成回复的token ID。
    generated_ids = model.generate(
        model_inputs.input_ids,
        # 生成最多512个新的tokens
        max_new_tokens=512
    )
    # 通过列表推导式，从generated_ids中排除输入部分的token ID，只保留生成的新内容。
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    """
    使用tokenizer.batch_decode将生成的token ID解码为文本。
    skip_special_tokens=True表示跳过特殊符号，使生成的文本更干净。
    """
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return response




# 加载原下载好的模型，Transformers加载模型权重（tokenizer和model）
tokenizer = AutoTokenizer.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", use_fast=False, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", device_map="auto", torch_dtype=torch.bfloat16)

# 加载训练好的Lora模型
model = PeftModel.from_pretrained(model, model_id="./output/Qwen2/checkpoint-400")

# 开始预测
test_texts = {
    'instruction': "你是一个文本分类领域的专家，你会接收到一段文本和几个潜在的分类选项，请输出文本内容的正确类型",
    'input': "文本:航空动力学报JOURNAL OF AEROSPACE POWER1998年 第4期 No.4 1998科技期刊管路系统敷设的并行工程模型研究*陈志英*　*　马　枚北京航空航天大学【摘要】　提出了一种应用于并行工程模型转换研究的标号法，该法是将现行串行设计过程(As-is)转换为并行设计过程(To-be)。本文应用该法将发动机外部管路系统敷设过程模型进行了串并行转换，应用并行工程过程重构的手段，得到了管路敷设并行过程模型。"
}

instruction = test_texts['instruction']
input_value = test_texts['input']

messages = [
    {"role": "system", "content": f"{instruction}"},
    {"role": "user", "content": f"{input_value}"}
]

response = predict(messages, model, tokenizer)
print(response)

